import random

from deap import base
from deap import creator
from deap import tools

import numpy as np
import matplotlib.pyplot as plt
from numpy.fft import fftshift,fft2,ifftshift,fftfreq,ifft2


N=25 #一个维度的采样数
d=1 #描述掩膜板尺寸
x=np.linspace(-0.1*d,1.1*d,N) 
y=x.copy()
X,Y=np.meshgrid(x,y)
Y=Y[::-1,:] #以上生成X,Y坐标网，尺寸略大于d以囊括衍射导致的偏大部分，对Y翻转，使其大小顺序符合习惯
mask=np.zeros((N,N))
mask[(X>0)&(Y>0)&(X<d)&(Y<d)&((X<0.2*d)|(Y>0.8*d)|((Y>0.4*d)&(Y<0.6*d)))]=1#生成F型掩膜板
mask=mask.reshape(N*N,1)#展成一维数组以便套用deap
kx = fftshift(fftfreq(N,1.2*d/N))
ky=kx.copy()
Kx,Ky=np.meshgrid(kx,ky)#得傅里叶变换结果数组中每个位置对应的波数
k = np.sqrt(Kx**2 + Ky**2)
NA=0.9#数值孔径
wavelength=2#波长 单位和d一样
fil = k <= 2*np.pi*NA /wavelength #确定可滤过的波数；同时尽量多的运算放在函数外，节省运行时间
dmask=mask.reshape(N*N,1)#dmask意为被优化后的掩膜

def image(mask):
    global X,Y,fil,N
    mask=np.array(mask)
    mask=mask.reshape(N,N)
    diffracted = fftshift(fft2(ifftshift(mask)))#焦平面处/动量空间像
    transmitted=diffracted*fil
    E=fftshift(ifft2(ifftshift(transmitted)))
    intensity=np.abs(E)**2
    return intensity/intensity.max()#返回归一化后的强度，即像；N*N

def dintensity(individual):
    absd=np.abs(image(individual).reshape(N*N,1)-mask)#absolute difference
    return sum(absd) #计算该掩膜所得图案与期望图案的差距


creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
creator.create("Individual", list, fitness=creator.FitnessMin)

toolbox = base.Toolbox()

# Attribute generator 
#                      define 'attr_bool' to be an attribute ('gene')
#                      which corresponds to integers sampled uniformly
#                      from the range [0,1] (i.e. 0 or 1 with equal
#                      probability)
toolbox.register("attr_bool", random.randint, 0, 1)

# Structure initializers
#                         define 'individual' to be an individual
#                         consisting of 100 'attr_bool' elements ('genes')
toolbox.register("individual", tools.initRepeat, creator.Individual, 
    toolbox.attr_bool, N**2)

# define the population to be a list of individuals
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

# the goal ('fitness') function to be maximized

#----------
# Operator registration
#----------
# register the goal / fitness function
toolbox.register("evaluate", dintensity)

# register the crossover operator

toolbox.register("mate", tools.cxTwoPoint)

# register a mutation operator with a probability to
# flip each attribute/gene of 0.05
toolbox.register("mutate", tools.mutFlipBit, indpb=0.05)

# operator for selecting individuals for breeding the next
# generation: each individual of the current generation
# is replaced by the 'fittest' (best) of three individuals
# drawn randomly from the current generation.
toolbox.register("select", tools.selTournament, tournsize=3)

#----------

def main():
    random.seed(64)

    # create an initial population of 300 individuals (where
    # each individual is a list of integers)
    pop = toolbox.population(n=300)

    # CXPB  is the probability with which two individuals
    #       are crossed
    #
    # MUTPB is the probability for mutating an individual
    CXPB, MUTPB = 0.5, 0.2

    print("Start of evolution")

    # Evaluate the entire population
    fitnesses = list(map(toolbox.evaluate, pop))
    for ind, fit in zip(pop, fitnesses):
        ind.fitness.values = fit

    print("  Evaluated %i individuals" % len(pop))

    # Extracting all the fitnesses of 
    fits = [ind.fitness.values[0] for ind in pop]

    # Variable keeping track of the number of generations
    g = 0

    # Begin the evolution
    while min(fits) >50  and g < 10000:
        # A new generation
        g = g + 1
        print("-- Generation %i --" % g)

        # Select the next generation individuals
        offspring = toolbox.select(pop, len(pop))
        # Clone the selected individuals
        offspring = list(map(toolbox.clone, offspring))

        # Apply crossover and mutation on the offspring
        for child1, child2 in zip(offspring[::2], offspring[1::2]):

            # cross two individuals with probability CXPB
            if random.random() < CXPB:
                toolbox.mate(child1, child2)

                # fitness values of the children
                # must be recalculated later
                del child1.fitness.values
                del child2.fitness.values

        for mutant in offspring:

            # mutate an individual with probability MUTPB
            if random.random() < MUTPB:
                toolbox.mutate(mutant)
                del mutant.fitness.values

        # Evaluate the individuals with an invalid fitness
        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        fitnesses = map(toolbox.evaluate, invalid_ind)
        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit

        print("  Evaluated %i individuals" % len(invalid_ind))

        # The population is entirely replaced by the offspring
        pop[:] = offspring

        # Gather all the fitnesses in one list and print the stats
        fits = [ind.fitness.values[0] for ind in pop]

        length = len(pop)
        mean = sum(fits) / length
        sum2 = sum(x*x for x in fits)
        std = abs(sum2 / length - mean**2)**0.5

        print("  Min %s" % min(fits))
        print("  Max %s" % max(fits))
        print("  Avg %s" % mean)
        print("  Std %s" % std)

    print("-- End of (successful) evolution --")

    best_ind = tools.selBest(pop, 1)[0]
    print("Best individual is %s, %s" % (best_ind, best_ind.fitness.values))

if __name__ == "__main__":
    main()
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        